{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "674efe61-226c-42b5-87d2-adfeee1e9cb0",
   "metadata": {},
   "source": [
    "# OpenAI Tool calling capabilities with Unity Catalog\n",
    "\n",
    "## Prerequisites\n",
    "\n",
    "**API Key**\n",
    "To run this tutorial, you will need an OpenAI API key. \n",
    "\n",
    "Once you have acquired your key, set it to the environment variable `OPENAI_API_KEY`.\n",
    "\n",
    "Below, we validate that this key is set properly in your environment.\n",
    "\n",
    "**Packages**\n",
    "\n",
    "To interface with both UnityCatalog and OpenAI, you will need to install the following packages:\n",
    "\n",
    "```shell\n",
    "pip install openai unitycatalog-openai\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f286c627-63d2-40bb-87a9-8ae68ddafc55",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "assert \"OPENAI_API_KEY\" in os.environ, (\n",
    "    \"Please set the OPENAI_API_KEY environment variable to your OpenAI API key\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8419dd3-619b-4b97-bd36-75ab72c60d6c",
   "metadata": {},
   "source": [
    "## Configuration and Client setup\n",
    "\n",
    "In order to connect to your Unity Catalog server, you'll need an instance of the `ApiClient` from the `unitycatalog-client` package. \n",
    "\n",
    "> Note: If you don't already have a Catalog and a Schema created, be sure to create them before running this notebook and adjust the `CATALOG` and `SCHEMA` variables below to suit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e9b8058-f8c2-489c-bb0f-8a93b8ce95a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from unitycatalog.ai.core.client import UnitycatalogFunctionClient\n",
    "from unitycatalog.ai.openai.toolkit import UCFunctionToolkit\n",
    "from unitycatalog.client import ApiClient, Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "627e4c12-6f46-44fe-b695-2dd1179b8dc1",
   "metadata": {},
   "outputs": [],
   "source": [
    "config = Configuration()\n",
    "config.host = \"http://localhost:8080/api/2.1/unity-catalog\"\n",
    "\n",
    "# The base ApiClient is async\n",
    "api_client = ApiClient(configuration=config)\n",
    "\n",
    "client = UnitycatalogFunctionClient(api_client=api_client)\n",
    "\n",
    "CATALOG = \"AICatalog\"\n",
    "SCHEMA = \"AISchema\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d4fc68b-55b8-4a6a-a5bc-cc1985bec728",
   "metadata": {},
   "source": [
    "## Define functions and register them to Unity Catalog\n",
    "\n",
    "In this next section, we'll be defining a Python function and creating it within Unity Catalog so that it can be retrieved and used as a tool with our calls to OpenAI. \n",
    "\n",
    "There are a few things to keep in mind when creating functions for use with the `create_python_function` API:\n",
    "\n",
    "- Ensure that your have properly defined types for all arguments and for the return of the function.\n",
    "- Ensure that you have a Google-style docstring defined that includes descriptions for the function, each argument, and the return of the function. This is critical, as these are used to populate the metadata associated with the function within Unity Catalog, providing contextual data for an LLM to understand when and how to call the tool associated with this function.\n",
    "- If there are packages being called that are not part of core Python, ensure that the import statements are locally scoped (defined within the function body)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "23bafac9-b256-4b87-9f2f-c9eb5060f7d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def impact_energy(mass: float, velocity: float, composition: str) -> float:\n",
    "    \"\"\"\n",
    "    Calculates the amount of energy transmitted to Earth by an interstellar body upon impact.\n",
    "\n",
    "    Args:\n",
    "        mass (float): Mass of the body in kilograms.\n",
    "        velocity (float): Velocity of the body in meters per second.\n",
    "        composition (str): Elemental composition of the body (e.g., 'iron', 'stone', 'ice', 'diamond', 'uranium', 'gold', 'titanium', 'lithium').\n",
    "\n",
    "    Returns:\n",
    "        Energy transmitted to Earth in joules.\n",
    "    \"\"\"\n",
    "    if mass <= 0:\n",
    "        raise ValueError(\"Mass must be positive.\")\n",
    "    if velocity <= 0:\n",
    "        raise ValueError(\"Velocity must be positive.\")\n",
    "\n",
    "    energy_transmission = {\n",
    "        \"iron\": 0.9,  # High survival rate through atmosphere\n",
    "        \"stone\": 0.5,  # Moderate survival rate\n",
    "        \"ice\": 0.1,  # Low survival rate due to ablation\n",
    "        \"diamond\": 0.8,  # High melting point and thermal conductivity\n",
    "        \"uranium\": 0.85,  # High density and melting point\n",
    "        \"gold\": 0.7,  # High density but lower melting point\n",
    "        \"titanium\": 0.75,  # High strength and melting point\n",
    "        \"lithium\": 0.2,  # Low density and melting point\n",
    "    }\n",
    "\n",
    "    transmission_coefficient = energy_transmission.get(composition.lower())\n",
    "\n",
    "    if transmission_coefficient is None:\n",
    "        raise ValueError(\n",
    "            f\"Unknown composition: {composition}. Must be one of: {list(energy_transmission.keys())}\"\n",
    "        )\n",
    "\n",
    "    kinetic_energy = 0.5 * mass * velocity**2\n",
    "\n",
    "    energy_transmitted = kinetic_energy * transmission_coefficient\n",
    "\n",
    "    return energy_transmitted"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "6f48677a-3f8c-4ebc-a813-516fd02c2436",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "FunctionInfo(name='impact_energy', catalog_name='AICatalog', schema_name='AISchema', input_params=FunctionParameterInfos(parameters=[FunctionParameterInfo(name='mass', type_text='DOUBLE', type_json='{\"name\": \"mass\", \"type\": \"double\", \"nullable\": false, \"metadata\": {\"comment\": \"Mass of the body in kilograms.\"}}', type_name=<ColumnTypeName.DOUBLE: 'DOUBLE'>, type_precision=None, type_scale=None, type_interval_type=None, position=0, parameter_mode=None, parameter_type=None, parameter_default=None, comment='Mass of the body in kilograms.'), FunctionParameterInfo(name='velocity', type_text='DOUBLE', type_json='{\"name\": \"velocity\", \"type\": \"double\", \"nullable\": false, \"metadata\": {\"comment\": \"Velocity of the body in meters per second.\"}}', type_name=<ColumnTypeName.DOUBLE: 'DOUBLE'>, type_precision=None, type_scale=None, type_interval_type=None, position=1, parameter_mode=None, parameter_type=None, parameter_default=None, comment='Velocity of the body in meters per second.'), FunctionParameterInfo(name='composition', type_text='STRING', type_json='{\"name\": \"composition\", \"type\": \"string\", \"nullable\": false, \"metadata\": {\"comment\": \"Elemental composition of the body (e.g., \\\\\"iron\\\\\", \\\\\"stone\\\\\", \\\\\"ice\\\\\", \\\\\"diamond\\\\\", \\\\\"uranium\\\\\", \\\\\"gold\\\\\", \\\\\"titanium\\\\\", \\\\\"lithium\\\\\").\"}}', type_name=<ColumnTypeName.STRING: 'STRING'>, type_precision=None, type_scale=None, type_interval_type=None, position=2, parameter_mode=None, parameter_type=None, parameter_default=None, comment='Elemental composition of the body (e.g., \"iron\", \"stone\", \"ice\", \"diamond\", \"uranium\", \"gold\", \"titanium\", \"lithium\").')]), data_type=<ColumnTypeName.DOUBLE: 'DOUBLE'>, full_data_type='DOUBLE', return_params=None, routine_body='EXTERNAL', routine_definition='if mass <= 0:\\n    raise ValueError(\"Mass must be positive.\")\\nif velocity <= 0:\\n    raise ValueError(\"Velocity must be positive.\")\\n\\nenergy_transmission = {\\n    \\'iron\\': 0.9,      # High survival rate through atmosphere\\n    \\'stone\\': 0.5,     # Moderate survival rate\\n    \\'ice\\': 0.1,       # Low survival rate due to ablation\\n    \\'diamond\\': 0.8,   # High melting point and thermal conductivity\\n    \\'uranium\\': 0.85,  # High density and melting point\\n    \\'gold\\': 0.7,      # High density but lower melting point\\n    \\'titanium\\': 0.75, # High strength and melting point\\n    \\'lithium\\': 0.2,   # Low density and melting point\\n}\\n\\ntransmission_coefficient = energy_transmission.get(composition.lower())\\n\\nif transmission_coefficient is None:\\n    raise ValueError(f\"Unknown composition: {composition}. Must be one of: {list(energy_transmission.keys())}\")\\n\\nkinetic_energy = 0.5 * mass * velocity ** 2\\n\\nenergy_transmitted = kinetic_energy * transmission_coefficient\\n\\nreturn energy_transmitted', routine_dependencies=None, parameter_style='S', is_deterministic=True, sql_data_access='NO_SQL', is_null_call=False, security_type='DEFINER', specific_name='impact_energy', comment='Calculates the amount of energy transmitted to Earth by an interstellar body upon impact.', properties='null', full_name='AICatalog.AISchema.impact_energy', owner=None, created_at=1732675815511, created_by=None, updated_at=1732675815511, updated_by=None, function_id='ae514cb7-cada-403f-b592-cdc0347e9627', external_language='PYTHON')"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "client.create_python_function(func=impact_energy, catalog=CATALOG, schema=SCHEMA, replace=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1890ac44-fa5d-4603-945e-5c26f41a4a07",
   "metadata": {},
   "source": [
    "## Build a tool calling Agent with OpenAI\n",
    "\n",
    "In order to let a GenAI service like a GPT model hosted by OpenAI use our functions, we need to register them as tools. \n",
    "\n",
    "To do this, we'll import the Unity Catalog AI OpenAI integration package and utilize the `UCFunctionToolkit` class to construct the interface we need to register tools. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3cf1675e-e935-40ef-9f17-b3ef7c9e5c1e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'type': 'function',\n",
       "  'function': {'name': 'AICatalog__AISchema__impact_energy',\n",
       "   'strict': True,\n",
       "   'parameters': {'properties': {'mass': {'description': 'Mass of the body in kilograms.',\n",
       "      'title': 'Mass',\n",
       "      'type': 'number'},\n",
       "     'velocity': {'description': 'Velocity of the body in meters per second.',\n",
       "      'title': 'Velocity',\n",
       "      'type': 'number'},\n",
       "     'composition': {'description': 'Elemental composition of the body (e.g., \"iron\", \"stone\", \"ice\", \"diamond\", \"uranium\", \"gold\", \"titanium\", \"lithium\").',\n",
       "      'title': 'Composition',\n",
       "      'type': 'string'}},\n",
       "    'title': 'AICatalog__AISchema__impact_energy__params',\n",
       "    'type': 'object',\n",
       "    'additionalProperties': False,\n",
       "    'required': ['mass', 'velocity', 'composition']},\n",
       "   'description': 'Calculates the amount of energy transmitted to Earth by an interstellar body upon impact.'}}]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "toolkit = UCFunctionToolkit(function_names=[f\"{CATALOG}.{SCHEMA}.impact_energy\"], client=client)\n",
    "tools = toolkit.tools\n",
    "tools"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "440854cb-5096-485a-a507-d5548f38caab",
   "metadata": {},
   "source": [
    "## Interface with OpenAI's GPT model with tool calling capabilities\n",
    "\n",
    "In the following code block, we submit our messages to OpenAI. As with a standard query, we provide both a system prompt message and a user message. \n",
    "\n",
    "In order to allow the LLM to utilize tool calling capabilities, we pass our toolkit definitions (`toolkit.tools`) to the `tools` argument within the OpenAI SDK `chat.completions.create` API. With the tool definitions provided, OpenAI's LLM is available to contextually 'decide' when it is appropriate to request a tool call to be executed with its response. \n",
    "\n",
    "In the case of the question that we're providing, OpenAI will recognize that it should call both of our tools in order to facilitate an accurate answer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b17969cd-b152-4efd-9910-324c557bbb3b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ChatCompletion(id='chatcmpl-AY2Q4daxucRRao7I4QFpqHSSqt4dW', choices=[Choice(finish_reason='tool_calls', index=0, logprobs=None, message=ChatCompletionMessage(content=None, refusal=None, role='assistant', audio=None, function_call=None, tool_calls=[ChatCompletionMessageToolCall(id='call_R9sXaB0d3oISGogH9awqCpFw', function=Function(arguments='{\"mass\":64008.5,\"velocity\":773542.99,\"composition\":\"iron\"}', name='AICatalog__AISchema__impact_energy'), type='function')]))], created=1732675816, model='gpt-4o-mini-2024-07-18', object='chat.completion', service_tier=None, system_fingerprint='fp_0705bf87c0', usage=CompletionUsage(completion_tokens=35, prompt_tokens=281, total_tokens=316, completion_tokens_details=CompletionTokensDetails(accepted_prediction_tokens=0, audio_tokens=0, reasoning_tokens=0, rejected_prediction_tokens=0), prompt_tokens_details=PromptTokensDetails(audio_tokens=0, cached_tokens=0)))"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import openai\n",
    "\n",
    "initial_messages = [\n",
    "    {\n",
    "        \"role\": \"system\",\n",
    "        \"content\": \"You are a helpful assistant that provides realistic, verbose, and hightly technical responses to abstract 'what if?' \"\n",
    "        \"questions that curious users ask. When responding to a question, utilize tools that are available to you to provide the most\"\n",
    "        \"accurate answers possible. After getting a response to a tool call, use that to explain in full detail the contextual \"\n",
    "        \"information to provide a full accounting of what the real-world effects would be for the question being posed and include \"\n",
    "        \"a thorough explanation that covers the topic in as much detail as you can.\",\n",
    "    },\n",
    "    {\n",
    "        \"role\": \"user\",\n",
    "        \"content\": \"What would happen if a sphere of iron with a mass of 64008.5kg struck the Earth traveling at 773542.99 m/s?\",\n",
    "    },\n",
    "]\n",
    "\n",
    "response = openai.chat.completions.create(\n",
    "    model=\"gpt-4o-mini\",\n",
    "    messages=initial_messages,\n",
    "    tools=tools,\n",
    ")\n",
    "\n",
    "response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "8624433f-01d9-41ad-8946-b74e3778c88a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'content': None,\n",
       "  'refusal': None,\n",
       "  'role': 'assistant',\n",
       "  'tool_calls': [{'id': 'call_R9sXaB0d3oISGogH9awqCpFw',\n",
       "    'function': {'arguments': '{\"mass\":64008.5,\"velocity\":773542.99,\"composition\":\"iron\"}',\n",
       "     'name': 'AICatalog__AISchema__impact_energy'},\n",
       "    'type': 'function'}]},\n",
       " {'role': 'tool',\n",
       "  'content': '{\"content\": \"1.7235308972987406e+16\"}',\n",
       "  'tool_call_id': 'call_R9sXaB0d3oISGogH9awqCpFw'}]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from unitycatalog.ai.openai.utils import generate_tool_call_messages\n",
    "\n",
    "messages = generate_tool_call_messages(response=response, client=client)\n",
    "\n",
    "messages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12fd494b-ebcb-4450-a328-771b8597e760",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "To understand the implications of a sphere of iron with a mass of 64,008.5 kg striking the Earth at a velocity of 773,542.99 m/s, we first need to calculate the energy released upon impact and then examine the potential consequences of such an event.\n",
      "\n",
      "### Energy Release Calculation\n",
      "The energy transmitted to Earth by this impact can be calculated using the formula for kinetic energy, which is given by:\n",
      "\n",
      "\\[\n",
      "E = \\frac{1}{2} m v^2\n",
      "\\]\n",
      "\n",
      "Where:\n",
      "- \\(E\\) is the kinetic energy,\n",
      "- \\(m\\) is the mass of the object (64,008.5 kg),\n",
      "- \\(v\\) is the velocity of the object (773,542.99 m/s).\n",
      "\n",
      "The calculation yields an impact energy of approximately \\(1.72 \\times 10^{16}\\) joules. To put this in perspective, this amount of energy is equivalent to about:\n",
      "\n",
      "- 4.1 Megatons of TNT.\n",
      "- Over four times the energy released by the atomic bomb dropped on Hiroshima, which was about 15 kilotons.\n",
      "\n",
      "### Potential Consequences of the Impact\n",
      "\n",
      "1. **Impact Location and Depth**: The exact effects would depend significantly on where the impact occurs. An impact into the ocean would result in massive tsunamis, while an impact on land would create an immediate crater and shockwave.\n",
      "\n",
      "2. **Crater Formation**: At this energy level, the explosion would create a substantial crater. For reference, the Chicxulub crater, which is associated with the dinosaur extinction event, was caused by an asteroid roughly 10 km in diameter. The energy from the iron sphere would create a significantly large crater, potentially kilometers in diameter.\n",
      "\n",
      "3. **Shockwave and Immediate Environment**: The shockwave generated by such an impact would obliterate everything within miles of the impact site due to the massive release of energy. In a terrestrial setting, this could level cities and create fires over large areas.\n",
      "\n",
      "4. **Global Effects**: The immediate effects of the impact would likely include massive dust and debris thrown into the atmosphere, potentially altering climate patterns temporarily. This fallout could block sunlight, leading to what is colloquially referred to as an \"impact winter\", which can have severe implications for global agriculture and ecosystems.\n",
      "\n",
      "5. **Seismic Activity**: The impact would also create significant seismic waves, potentially triggering earthquakes and volcanic activity in the surrounding region.\n",
      "\n",
      "6. **Long-term Effects**: Should the impact create a large enough dust cloud, this could lead to a significant drop in global temperatures for an extended period, affecting weather patterns and potentially leading to food shortages.\n",
      "\n",
      "In summary, an iron sphere of 64,008.5 kg striking Earth at such a high velocity would unleash catastrophic forces capable of causing extensive localized destruction and potentially global climatic changes.\n"
     ]
    }
   ],
   "source": [
    "# Prepend the message history\n",
    "final_messages = initial_messages + messages\n",
    "\n",
    "response = openai.chat.completions.create(\n",
    "    model=\"gpt-4o-mini\",\n",
    "    messages=final_messages,\n",
    "    tools=tools,\n",
    ")\n",
    "\n",
    "response.choices[0].message.content"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
